package com.commonEncryptionAlgorithms.symmetry.AES;

import com.util.Base64Util;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import javax.crypto.spec.IvParameterSpec;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

public class AesEncryption {
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";
    private static final String SPLIT_STR = "|";
    private static final int IV_LENGTH = 16;

    /**
     * 生成key，根据秘钥和密码 长度
     */
    public static SecretKey generateKey(String password, int passwordLen) {
        try {
            // 使用密码生成密钥
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] keyBytes = digest.digest(password.getBytes(StandardCharsets.UTF_8));

            // 由于SHA-256生成的散列值长度为32字节，我们需要将其截断到所需的密钥长度
            byte[] truncatedKeyBytes = new byte[passwordLen / 8];
            System.arraycopy(keyBytes, 0, truncatedKeyBytes, 0, passwordLen / 8);

            // 使用截断后的字节序列创建SecretKeySpec
            return new SecretKeySpec(truncatedKeyBytes, ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            return null;
        }
    }
    //TODO: 可以有多个生成key的方式


    /**
     * 生成一个随机的向量坐标iv
     *
     * @return 返回iv的byte数组
     */
    private static byte[] generateIv() {
        byte[] iv = new byte[IV_LENGTH];
        new SecureRandom().nextBytes(iv);
        return iv;
    }

    /**
     * 加密
     *
     * @param strToEncrypt 加密字符串
     * @return 返回加密后的字符串，是一个Base64字符串
     */
    public static String encipher(String strToEncrypt, SecretKey secretKey) throws Exception {
        byte[] ivByte = generateIv();
        byte[] encryptByte = encipher(strToEncrypt, secretKey, ivByte);
        String ivStr = Base64Util.encode(ivByte);
        String encodeStr = Base64Util.encode(encryptByte);
        return ivStr + SPLIT_STR + encodeStr;
    }

    /**
     * 加密
     *
     * @param bytesToEncrypt 加密byte数组
     * @return 返回一个加密后的byte数组，这个数组是base64 编码之后的byte数组。 一般这个用于转换为文件。
     */
    public static byte[] encipher(byte[] bytesToEncrypt, SecretKey secretKey) throws Exception {
        byte[] ivByte = generateIv();
        byte[] encryptByte = encipher(bytesToEncrypt, secretKey, ivByte);
        String ivStr = Base64Util.encode(ivByte);
        String encodeStr = Base64Util.encode(encryptByte);
        String str = ivStr + SPLIT_STR + encodeStr;
        // 再对拼接之后的字符串进行一次加密。
        String encipher = encipher(str, secretKey);
        return encipher.getBytes(StandardCharsets.UTF_8);
    }

    /**
     * 加密
     *
     * @param bytes     原始 byte 数组
     * @param secretKey key对象
     * @param iv        iv坐标
     * @return 返回加密后的byte数组
     */
    private static byte[] encipher(byte[] bytes, SecretKey secretKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivParameterSpec);
        return cipher.doFinal(bytes);
    }

    /**
     * 加密
     *
     * @param strToEncrypt 原始 字符串
     * @param secretKey    key对象
     * @param iv           iv坐标
     * @return 返回加密后的byte数组
     */
    private static byte[] encipher(String strToEncrypt, SecretKey secretKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivParameterSpec);
        return cipher.doFinal(strToEncrypt.getBytes());
    }

    /**
     * 解密
     *
     * @param strToDecrypt 需要解密的字符串
     * @return 返回解密结果
     */
    public static String decrypt(String strToDecrypt, SecretKey secretKey) throws Exception {
        String[] split = strToDecrypt.split("\\|");
        String ivStr = split[0];
        byte[] ivByte = Base64Util.decode(ivStr);
        String encryptStr = split[1];
        byte[] encryptBytes = Base64Util.decode(encryptStr);
        return decrypt(encryptBytes, secretKey, ivByte);
    }

    public static byte[] decrypt(byte[] bytes, SecretKey secretKey) throws Exception {
        // 首先转换为 字符
        String encode = new String(bytes, StandardCharsets.UTF_8);
        // 然后进行一个解码
        String decryptStr = decrypt(encode, secretKey);
        String[] split = decryptStr.split("\\|");
        String ivStr = split[0];
        byte[] ivByte = Base64Util.decode(ivStr);
        String encryptStr = split[1];
        byte[] encryptBytes = Base64Util.decode(encryptStr);
        return decryptToByte(encryptBytes, secretKey, ivByte);
    }

    /**
     * 解密
     *
     * @param strToDecrypt 解密字符串
     * @param secretKey    解密key对象
     * @param iv           iv坐标
     * @return 返回解密后的字符串
     */
    private static String decrypt(byte[] strToDecrypt, SecretKey secretKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivParameterSpec);
        byte[] decryptedBytes = cipher.doFinal(strToDecrypt);
        return new String(decryptedBytes);
    }
    /**
     * 解密
     *
     * @param strToDecrypt 解密字符串
     * @param secretKey    解密key对象
     * @param iv           iv坐标
     * @return 返回解密后的字符串
     */
    private static byte[] decryptToByte(byte[] strToDecrypt, SecretKey secretKey, byte[] iv) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivParameterSpec);
        return cipher.doFinal(strToDecrypt);
    }
}